#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识库逐行搜索工具
功能：在逐行向量化的知识库中搜索相关内容
"""

from pymilvus import connections, Collection
import re
import jieba
from sklearn.feature_extraction.text import TfidfVectorizer


def connect_to_milvus(host="localhost", port="19530"):
    """连接到Milvus"""
    try:
        connections.connect("default", host=host, port=port)
        print(f"✅ 成功连接到Milvus服务器 {host}:{port}")
        return True
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return False


def search_by_keyword(collection_name, keyword, limit=10):
    """按关键词搜索"""
    try:
        col = Collection(collection_name)
        col.load()
        
        # 使用like查询搜索包含关键词的行
        expr = f'line_text like "%{keyword}%"'
        results = col.query(
            expr=expr,
            limit=limit,
            output_fields=['filename', 'line_number', 'line_text']
        )
        
        print(f"\n🔍 搜索关键词: '{keyword}'")
        print(f"找到 {len(results)} 条相关记录:")
        
        for i, result in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(f"文件: {result['filename']}")
            print(f"行号: {result['line_number']}")
            
            # 高亮显示关键词
            text = result['line_text']
            highlighted = text.replace(keyword, f"【{keyword}】")
            print(f"内容: {highlighted}")
        
        return results
        
    except Exception as e:
        print(f"❌ 搜索失败: {e}")
        return []


def search_by_file(collection_name, filename_pattern, limit=20):
    """按文件名搜索"""
    try:
        col = Collection(collection_name)
        col.load()
        
        expr = f'filename like "{filename_pattern}%"'
        results = col.query(
            expr=expr,
            limit=limit,
            output_fields=['filename', 'line_number', 'line_text']
        )
        
        print(f"\n📁 搜索文件: '{filename_pattern}*'")
        print(f"找到 {len(results)} 行记录:")
        
        for i, result in enumerate(results, 1):
            print(f"{i}. 行{result['line_number']}: {result['line_text'][:80]}...")
        
        return results
        
    except Exception as e:
        print(f"❌ 搜索失败: {e}")
        return []


def get_collection_stats(collection_name):
    """获取集合统计信息"""
    try:
        col = Collection(collection_name)
        col.load()
        
        print(f"\n📊 集合统计: {collection_name}")
        print(f"总记录数: {col.num_entities}")
        
        # 获取文件统计
        all_results = col.query(
            expr='line_number > 0',
            output_fields=['filename']
        )
        
        file_counts = {}
        for result in all_results:
            filename = result['filename']
            file_counts[filename] = file_counts.get(filename, 0) + 1
        
        print(f"文件数量: {len(file_counts)}")
        print(f"平均每文件行数: {col.num_entities / len(file_counts):.1f}")
        
        # 显示行数最多的文件
        top_files = sorted(file_counts.items(), key=lambda x: x[1], reverse=True)[:5]
        print("\n📈 行数最多的文件:")
        for filename, count in top_files:
            print(f"  {filename}: {count} 行")
        
        return file_counts
        
    except Exception as e:
        print(f"❌ 获取统计失败: {e}")
        return {}


def interactive_search():
    """交互式搜索"""
    collection_name = "knowledge_lines"
    
    print("🔍 知识库逐行搜索工具")
    print("=" * 50)
    
    if not connect_to_milvus():
        return
    
    # 显示集合统计
    get_collection_stats(collection_name)
    
    while True:
        print("\n" + "=" * 50)
        print("请选择搜索方式:")
        print("1. 关键词搜索")
        print("2. 文件名搜索")
        print("3. 查看集合统计")
        print("4. 退出")
        
        choice = input("请输入选择 (1-4): ").strip()
        
        if choice == "1":
            keyword = input("请输入搜索关键词: ").strip()
            if keyword:
                try:
                    limit = int(input("返回结果数量 (默认10): ") or "10")
                except ValueError:
                    limit = 10
                search_by_keyword(collection_name, keyword, limit)
        
        elif choice == "2":
            pattern = input("请输入文件名前缀 (如 'C1' 查找C1开头的文件): ").strip()
            if pattern:
                try:
                    limit = int(input("返回结果数量 (默认20): ") or "20")
                except ValueError:
                    limit = 20
                search_by_file(collection_name, pattern, limit)
        
        elif choice == "3":
            get_collection_stats(collection_name)
        
        elif choice == "4":
            print("👋 再见!")
            break
        
        else:
            print("❌ 无效选择")


def demo_searches():
    """演示搜索"""
    collection_name = "knowledge_lines"
    
    print("🎯 演示搜索功能")
    print("=" * 30)
    
    if not connect_to_milvus():
        return
    
    # 演示关键词搜索
    print("\n1. 搜索包含'模板'的行:")
    search_by_keyword(collection_name, "模板", 3)
    
    print("\n2. 搜索包含'PCR'的行:")
    search_by_keyword(collection_name, "PCR", 3)
    
    print("\n3. 查看C113文件的内容:")
    search_by_file(collection_name, "C113", 5)


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "demo":
        demo_searches()
    else:
        interactive_search()
